from core import *
import matplotlib.pyplot as plt
import numpy as np
from visualizer.basic import *
from sklearn.metrics import mean_squared_error


def Ts(ts2, ts1, l, m, step, func):
	return TE(ts2, ts1, l, m, step, func) - TE(ts1, ts2, l, m, step, func)


def extremum_1(subsequence):
	return array2index([np.argmin(subsequence), np.argmax(subsequence)])


def extremum_2(subsequence):
	if len(subsequence) <= 2:
		return array2index(np.argsort(subsequence))
	else:
		array = np.argsort(subsequence)
		array = np.append(array[:2], array[-2:])
		return array2index(array)


def extremum_3(subsequence):
	if len(subsequence) <= 3:
		return array2index(np.argsort(subsequence))
	else:
		array = np.argsort(subsequence)
		array = np.append(array[:3], array[-3:])
		return array2index(array)


if __name__ == '__main__':
	ts1 = np.loadtxt('../data/pr_NY.csv')  # 降水量 # Y
	ts2 = np.loadtxt('../data/tas_NY.csv')  # 温度 # X
	N = 21
	l = 1
	m = 3
	step = 1
	symbol_0 = []
	symbol_2 = []
	symbol_4 = []
	symbol_6 = []
	# symbol
	x = list(range(3, N))
	for m in range(3, N):
		T = Ts(ts2, ts1, l, m, step, symbol)
		symbol_0.append(T)
	plt.plot(x, symbol_0, label='STE')
	# --------------------------------------------------------------------------------
	# extremum 1
	ex_3 = []
	for m in range(3, N):
		T = Ts(ts2, ts1, l, m, step, extremum_3)
		ex_3.append(T)
	plt.plot(x, ex_3, label='$t=3$')
	# extremum 2
	ex_2 = []
	for m in range(3, N):
		T = Ts(ts2, ts1, l, m, step, extremum_2)
		ex_2.append(T)
	plt.plot(x, ex_2, label='$t=2$')
	# extreme 1
	ex_1 = []
	for m in range(3, N):
		T = Ts(ts2, ts1, l, m, step, extremum_1)
		ex_1.append(T)
	plt.plot(x, ex_1, label='$t=1$')
	# --------------------------------------------------------------------------------
	plt.xlabel('$m$')
	plt.ylabel('$T(m)$')
	plt.xticks(np.arange(3, N, 3), np.arange(3, N, 3))
	plt.legend()
	plt.tight_layout()
	plt.savefig(PATH+'TE.pdf', format='pdf')
	plt.show()
	mse1 = mean_squared_error(ex_1, symbol_0)
	mse2 = mean_squared_error(ex_2, symbol_0)
	mse3 = mean_squared_error(ex_3, symbol_0)
	print('MSE($t$=1):', mse1)
	print('MSE($t$=2):', mse2)
	print('MSE($t$=3):', mse3)
	